use async_trait::async_trait;
use log::{debug, info};
use pingora::prelude::*;
use std::sync::Arc;
use std::time::Duration;
use pingora::lb::health_check::HttpHealthCheck;
// 添加 CryptoProvider 相关引入
use rustls::crypto::CryptoProvider;
use once_cell::sync::Lazy;
use pingora_http::{RequestHeader, ResponseHeader};
use pingora_limits::rate::Rate;
use pingora_load_balancing::prelude::{RoundRobin, TcpHealthCheck};
use pingora_load_balancing::LoadBalancer;
use pingora_proxy::{http_proxy_service, ProxyHttp, Session as ProxyHttpSession};
use log::error;
use pingora_core::server::configuration::{Opt as ServerOpt, ServerConf};
use pingora_core::server::Server;
use module_utils::{FromYaml, merge_opt, merge_conf, RequestFilter};
use serde::Deserialize;
use std::fs::File;
use std::io::BufReader;
//use static_files_module::{CompressionAlgorithm, StaticFilesConf, StaticFilesHandler};
use structopt::StructOpt;
//https://docs.rs/static-files-module/latest/static_files_module/

// #[merge_opt]
// struct MyServerOpt {
//     server: ServerOpt,
//     static_files: StaticFilesOpt,
// }



// The configuration settings from both structures are merged, so that the user doesn't need to
// care which structure they belong to.
// #[merge_conf]
// struct MyServerConf {
//     server: ServerConf,
//     static_files: StaticFilesConf,
// }

// The command line flags from both structures are merged, so that the user doesn't need to
// care which structure they belong to.

pub struct MyWebProxyServer {
    upstreams: Arc<LoadBalancer<RoundRobin>>,
   // static_files_handler: StaticFilesHandler,
}

impl MyWebProxyServer {
    pub fn new(upstreams: Arc<LoadBalancer<RoundRobin>>) -> Self {
        Self {
            upstreams,
        }
    }
    // pub fn new(upstreams: Arc<LoadBalancer<RoundRobin>>, static_files_handler: StaticFilesHandler) -> Self {
    //     Self {
    //         upstreams,
    //         static_files_handler,
    //     }
    // }
    //
    pub fn get_request_appid(&self, session: &mut Session) -> Option<String> {
        match session
            .req_header()
            .headers
            .get("appid")
            .map(|v| v.to_str())
        {
            None => None,
            Some(v) => match v {
                Ok(v) => Some(v.to_string()),
                Err(_) => None,
            },
        }
    }
}

// Rate limiter
static RATE_LIMITER: Lazy<Rate> = Lazy::new(|| Rate::new(Duration::from_secs(1)));

// max request per second per client
static MAX_REQ_PER_SEC: isize = 1;
#[async_trait]
impl ProxyHttp for MyWebProxyServer {
    /// For this small example, we don't need context storage
    type CTX = ();
    fn new_ctx(&self) -> () {
        ()
    }
    // type CTX = <StaticFilesHandler as RequestFilter>::CTX;
    // fn new_ctx(&self) -> Self::CTX {
    //     StaticFilesHandler::new_ctx()
    // }

    async fn upstream_peer(&self, _session: &mut ProxyHttpSession, _ctx: &mut ()) -> Result<Box<HttpPeer>> {
        let upstream = self
            .upstreams
            .select(b"", 256) // hash doesn't matter for round robin
            .unwrap();

        debug!("upstream peer is: {}", upstream.to_string());

        // Set SNI to one.one.one.one
        let peer = Box::new(HttpPeer::new(
            upstream,
            false,
            "one.one.one.one".to_string(),
        ));
        Ok(peer)
    }
    async fn request_filter(&self, session: &mut ProxyHttpSession, _ctx: &mut Self::CTX) -> Result<bool>
    where
        Self::CTX: Send + Sync,
    {
        //session.downstream_compression.adjust_level(3);
        // 获得path
        let path = session.req_header().uri.path();
        debug!("path: {}", path);
        if path.eq("/") {
            //self.static_files_handler.handle(session, _ctx).await;
            return Ok(true);
        }else {
            let appid = match self.get_request_appid(session) {
                None => return Ok(false), // no client appid found, skip rate limiting
                Some(addr) => addr,
            };

            // retrieve the current window requests
            let curr_window_requests = RATE_LIMITER.observe(&appid, 1);
            if curr_window_requests > MAX_REQ_PER_SEC {
                // rate limited, return 429
                let mut header = ResponseHeader::build(429, None).unwrap();
                header
                    .insert_header("X-Rate-Limit-Limit", MAX_REQ_PER_SEC.to_string())
                    .unwrap();
                header.insert_header("X-Rate-Limit-Remaining", "0").unwrap();
                header.insert_header("X-Rate-Limit-Reset", "1").unwrap();
                session.set_keepalive(None);
                session
                    .write_response_header(Box::new(header), true)
                    .await?;
                return Ok(true);
            }
            Ok(false)
        }
    }
    async fn upstream_request_filter(
        &self,
        _session: &mut Session,
        upstream_request: &mut RequestHeader,
        _ctx: &mut Self::CTX,
    ) -> Result<()> {
        upstream_request
            .insert_header("Host", "one.one.one.one")?;
        upstream_request
            .insert_header("X-Forwarded-For", "127.0.0.1")?;
        upstream_request
            .insert_header("User-Agent", "PingCAP-Proxy/1.0")?;
        Ok(())
    }
}
fn main() {
    // 初始化 CryptoProvider，使用 ring 作为默认提供者
    CryptoProvider::install_default(rustls::crypto::ring::default_provider())
        .expect("Failed to install default crypto provider");

    println!("Hello, world!");
    env_logger::builder()
        .filter_level(log::LevelFilter::Info)
        .filter_module("web0007", log::LevelFilter::Trace)
        .init();

    let server_conf = ServerConf::load_from_yaml("conf.yaml").ok().unwrap();
    info!("server_conf: {:?}", server_conf);

    let opt = ServerOpt::parse_args() ;
    info!("opt: {:?}", opt);

     let mut my_server = Server::new_with_opt_and_conf(opt, server_conf);

    //let mut my_server = Server::new(None).unwrap();
    my_server.bootstrap();

    let mut upstreams = LoadBalancer::try_from_iter(["127.0.0.1:8080", "127.0.0.1:8081"]).unwrap();

    let hc = TcpHealthCheck::new();
    let hc_http_8080 = Box::new(HttpHealthCheck::new("127.0.0.1:8080",false));
    let hc_http_8081 = Box::new(HttpHealthCheck::new("127.0.0.1:8081",false));
    upstreams.set_health_check(hc_http_8080);
    upstreams.set_health_check(hc_http_8081);

    upstreams.set_health_check(hc);
    upstreams.health_check_frequency = Some(Duration::from_secs(10));
    // Set background service
    let background = background_service("health check", upstreams);
    let upstreams = background.task();
    //
    // let static_conf = StaticFilesConf {
    //     root: Some("dist".into()),
    //     precompressed: vec![CompressionAlgorithm::Gzip, CompressionAlgorithm::Brotli],
    //     index_file: vec!["index.html".to_string()],
    //     page_404: Some("index.html".to_string()),
    //     ..Default::default()
    // };
    // info!("static_conf: {:?}", static_conf);
    // let static_files_handler: StaticFilesHandler = static_conf.try_into().unwrap();

    let mut lb = http_proxy_service(&my_server.configuration, MyWebProxyServer::new(upstreams));
    lb.add_tcp("0.0.0.0:6288");

    my_server.add_service(lb);
    my_server.add_service(background);
    my_server.run_forever();
}

